import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import time
import requests
import zl_money_flow as zl
import stock_zt as zt

# ----------------------
# 1. 基础功能（不变）
# ----------------------
def is_trading_day(date):
    weekday = date.weekday()
    return weekday < 5


def get_current_period():
    now = datetime.now()
    today = now.date()
    if not is_trading_day(today):
        return "非交易日"
    morning_start = datetime(today.year, today.month, today.day, 9, 30)
    morning_end = datetime(today.year, today.month, today.day, 11, 30)
    afternoon_start = datetime(today.year, today.month, today.day, 13, 0)
    afternoon_end = datetime(today.year, today.month, today.day, 15, 0)
    if now < morning_start:
        return "盘前"
    elif morning_start <= now <= morning_end:
        return "上午交易中"
    elif morning_end < now < afternoon_start:
        return "午休（上午已收盘）"
    elif afternoon_start <= now <= afternoon_end:
        return "下午交易中"
    else:
        return "盘后"


def get_target_date():
    now = datetime.now()
    period = get_current_period()
    if period in ["盘前", "非交易日"]:
        target_date = (now - timedelta(days=1)).date()
        while not is_trading_day(target_date):
            target_date -= timedelta(days=1)
        return target_date
    else:
        return now.date()


# ----------------------
# 2. 实时价格获取（不变）
# ----------------------
def get_sina_latest_price(stock_code):
    if stock_code.startswith('6'):
        full_code = f"sh{stock_code}"
    else:
        full_code = f"sz{stock_code}"
    url = f"http://hq.sinajs.cn/list={full_code}"
    try:
        response = requests.get(url, timeout=5)
        if response.status_code == 200:
            data_str = response.text.split('=')[1].strip('";\n')
            data_list = data_str.split(',')
            if len(data_list) >= 4:
                current_price = float(data_list[3])
                prev_close = float(data_list[2])
                return {
                    "代码": stock_code,
                    "最新价格": current_price,
                    "昨日收盘价": prev_close,
                    "数据类型": "最新价"
                }
            else:
                print(f"⚠️ {stock_code}数据字段不足")
    except Exception as e:
        print(f"⚠️ {stock_code}获取失败：{e}")
    return None


# ----------------------
# 3. 无限制分析（以首日主力净流入为基准）
# ----------------------
def analyze_stocks_without_filter(hist_df):
    """无限制条件分析：以首日主力净流入为基准计算资金留存率"""
    print(f"\n【分析所有股票（无限制条件）】")
    total_stocks = hist_df["代码"].nunique()
    print(f"总股票数：{total_stocks}")

    stock_metrics = []
    for code, group in hist_df.groupby("代码"):
        stock_name = group["名称"].iloc[0]
        group_sorted = group.sort_values("日期")

        # 取首日数据作为基准（不再要求涨停）
        first_day = group_sorted["日期"].min()
        first_day_data = group_sorted[group_sorted["日期"] == first_day]
        base_fund = first_day_data["主力净流入-净额（万元）"].iloc[0]
        base_close = first_day_data["收盘价"].iloc[0]

        # 计算资金留存率（以首日主力净流入为基准值1）
        group_sorted["相对基准资金"] = group_sorted["主力净流入-净额（万元）"] / base_fund
        group_sorted["累计相对资金"] = group_sorted["相对基准资金"].cumsum()
        group_sorted["资金留存率(%)"] = group_sorted["累计相对资金"] * 100

        # 计算股价相对涨幅
        group_sorted["股价相对首日(%)"] = (group_sorted["收盘价"] / base_close) * 100
        group_sorted["股价涨幅(%)"] = group_sorted["股价相对首日(%)"] - 100

        # 取最近一个交易日的数据
        latest_data = group_sorted.iloc[-1]

        # 核心指标：资金留存率 - 股价相对值
        fund_price_diff = latest_data["资金留存率(%)"] - latest_data["股价相对首日(%)"]

        # 保存结果
        stock_metrics.append({
            "代码": code,
            "名称": stock_name,
            "基准日": first_day.strftime("%Y-%m-%d"),
            "基准日主力净流入(万元)": base_fund,
            "资金留存率(%)": round(latest_data["资金留存率(%)"], 2),
            "股价涨幅(%)": round(latest_data["股价涨幅(%)"], 2),
            "资金股价差值(%)": round(fund_price_diff, 2)
        })

    # 按资金股价差值降序排序
    return sorted(stock_metrics, key=lambda x: x["资金股价差值(%)"], reverse=True)


# ----------------------
# 4. 主逻辑（调用无限制分析函数）
# ----------------------
def main():
    try:
        # 加载数据
        print("【1/4】加载历史数据...")
        hist_df = zl.all_zt_after_data.copy()
        hist_df["代码"] = hist_df["代码"].apply(lambda x: str(x).zfill(6))
        hist_df["日期"] = pd.to_datetime(hist_df["日期"]).dt.date
        print(f"历史数据加载完成（股票数：{hist_df['代码'].nunique()}，记录数：{len(hist_df)}）")

        # 分析股票（无限制条件）
        print("\n【2/4】分析股票资金留存率...")
        analyzed_stocks = analyze_stocks_without_filter(hist_df)
        if not analyzed_stocks:
            print("❌ 无有效数据")
            return

        # 获取最新价格
        print("\n【3/4】获取最新价格...")
        target_date = get_target_date()
        current_period = get_current_period()
        print(f"当前时段：{current_period}，目标日期：{target_date}")

        all_codes = [stock["代码"] for stock in analyzed_stocks]
        name_map = {stock["代码"]: stock["名称"] for stock in analyzed_stocks}

        # 获取价格数据
        price_data_list = []
        if current_period in ["上午交易中", "下午交易中", "午休（上午已收盘）"]:
            for code in all_codes:
                print(f"获取 {code}（{name_map[code]}）实时价格...")
                price_data = get_sina_latest_price(code)
                if price_data:
                    price_data_list.append({
                        "代码": code,
                        "当前价格": price_data["最新价格"],
                        "基准价格": price_data["昨日收盘价"],
                        "数据类型": "实时价格"
                    })
                time.sleep(0.3)
        else:
            print("非交易时段，获取收盘价格...")
            for code in all_codes:
                code_hist = hist_df[hist_df["代码"] == code]
                target_data = code_hist[code_hist["日期"] == target_date]
                if target_data.empty:
                    latest_date = max(code_hist["日期"])
                    target_data = code_hist[code_hist["日期"] == latest_date]
                    print(f"⚠️ {code}使用最近日期{latest_date}数据")
                if not target_data.empty:
                    close_price = target_data["收盘价"].iloc[0]
                    prev_dates = code_hist[code_hist["日期"] < target_data["日期"].iloc[0]]["日期"]
                    prev_close = code_hist[code_hist["日期"] == max(prev_dates)]["收盘价"].iloc[
                        0] if not prev_dates.empty else close_price
                    price_data_list.append({
                        "代码": code,
                        "当前价格": close_price,
                        "基准价格": prev_close,
                        "数据类型": "收盘价格"
                    })

        if not price_data_list:
            print("❌ 未获取到价格数据")
            return

        # 整理结果
        print("\n【4/4】整理结果...")
        result_list = []
        for price_data in price_data_list:
            code = price_data["代码"]
            stock_info = next(s for s in analyzed_stocks if s["代码"] == code)
            current_price = price_data["当前价格"]
            base_price = price_data["基准价格"]
            daily_change = round((current_price - base_price) / base_price * 100, 2)

            result_list.append({
                "股票代码": code,
                "股票名称": name_map[code],
                "基准日": stock_info["基准日"],
                "基准日主力净流入(万元)": stock_info["基准日主力净流入(万元)"],
                "资金留存率(%)": stock_info["资金留存率(%)"],
                "股价涨幅(%)": stock_info["股价涨幅(%)"],
                "资金股价差值(%)": stock_info["资金股价差值(%)"],
                "基准价格(元)": round(base_price, 2),
                "最新价格(元)": round(current_price, 2),
                "当日涨幅(%)": daily_change,
                "数据时间": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            })

        # 保存并输出
        result_df = pd.DataFrame(result_list)
        output_path = zt.get_file_path("结果排名.csv")
        result_df.to_csv(output_path, index=False, encoding="utf-8-sig")
        print(f"\n✅ 结果已保存至：{output_path}")
        print("\n【分析结果（前5条）】")
        print(result_df.head().to_string(index=False))

    except Exception as e:
        print(f"\n❌ 程序出错：{e}")


if __name__ == "__main__":
    main()